import numpy as np
from scipy import stats
import matplotlib.pyplot as plt

def binom_pmf(n=1, p=0.1):
    """
    二项分布有两个参数
    https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.binom.html#scipy.stats.binom
    :param n:试验次数
    :param p:单次实验成功的概率
    :return:
    """
    binom_dis = stats.binom(n, p)
    x = np.arange(binom_dis.ppf(0.0001), binom_dis.ppf(0.9999))
    print(x)  # [ 0.  1.  2.  3.  4.]
    fig, ax = plt.subplots(1, 1)
    ax.plot(x, binom_dis.pmf(x), 'bo', label='binom pmf')
    ax.vlines(x, 0, binom_dis.pmf(x), colors='b', lw=5, alpha=0.5)
    ax.legend(loc='best', frameon=False)
    plt.ylabel('Probability')
    plt.title('PMF of binomial distribution(n={}, p={})'.format(n, p))
    plt.show()

def binom_dis(n=1, p=0.1):
    """
    二项分布，模拟抛硬币试验
    :param n: 实验总次数
    :param p: 单次实验成功的概率
    :return: 试验成功的次数
    """
    binom_dis = stats.binom(n, p)
    simulation_result = binom_dis.rvs(size=5)  # 取5个符合该分布的随机变量
    print(simulation_result)  # [ 7 11 13  8 13], 每次结果会不一样
    prob_10 = binom_dis.pmf(10)
    print(prob_10)  # 0.117

binom_dis(n=20, p=0.6)
binom_pmf(n=20, p=0.6)